#include "memory_allocator_impl.hh"
#include "../../thirdparty/klogger/klogger/interface/logger.h"
#include "../detail/box_alloc.hh"
#include "../util/time_util.hh"
#include "../util/string_util.hh"
#include <cstdlib>
#include <sstream>
#include <typeinfo>

kratos::service::MemoryAllocatorImpl::MemoryAllocatorImpl() {}

kratos::service::MemoryAllocatorImpl::~MemoryAllocatorImpl() {}

auto kratos::service::MemoryAllocatorImpl::alloc(std::size_t size) -> void * {
  return reinterpret_cast<void *>(box_malloc(size));
}

auto kratos::service::MemoryAllocatorImpl::dealloc(void *ptr) -> void {
#if defined(DEBUG) || defined(_DEBUG)
  auto mem_info_it = current_snapshot_.mem_info_map_.find(
      reinterpret_cast<std::ptrdiff_t>(ptr));
  if (mem_info_it != current_snapshot_.mem_info_map_.end()) {
    auto hash_name_it = current_snapshot_.hash_name_map_.find(
        mem_info_it->second.count_info_key);
    if (hash_name_it != current_snapshot_.hash_name_map_.end()) {
      hash_name_it->second.count -= 1;
    }
  }
#endif // defined(DEBUG) || defined(_DEBUG)
  box_free(ptr);
}

auto kratos::service::MemoryAllocatorImpl::alloc(
  std::size_t size,
  std::size_t hash_id,
  const std::string &name,
  const std::string &file_name,
  int line) -> void * {
#if defined(DEBUG) || defined(_DEBUG)
  auto *addr = alloc(size);
  if (!addr) {
    return nullptr;
  }
  std::size_t key = hash_id;
  auto it = current_snapshot_.hash_name_map_.find(hash_id);
  if (it != current_snapshot_.hash_name_map_.end()) {
    it->second.count += 1;
    it->second.history_count += 1;
  } else {
    current_snapshot_.hash_name_map_[hash_id] = {1, util::demangle(name), 1};
  }
  auto ptr_key = reinterpret_cast<std::ptrdiff_t>(addr);
  current_snapshot_.mem_info_map_[ptr_key] = {file_name, line, hash_id};
  return addr;
#else
  return alloc(size);
#endif // defined(DEBUG) || defined(_DEBUG)
}

auto kratos::service::MemoryAllocatorImpl::print_snapshot(
    const std::string &name, std::ostream &os, const MemorySnapshot &snapshot,
    std::time_t ts) -> void {
#if defined(DEBUG) || defined(_DEBUG)
  os << std::endl
     << "  >============================" + name + "============================<"
     << std::endl
     << "  | Timestamp:" << ts << std::endl;
  for (const auto &it : snapshot.hash_name_map_) {
    os << "  +----------------------------------------------------------"
       << std::endl;
    os << "  | Type:" << it.second.name << std::endl;
    os << "  | Current Count:" << it.second.count << std::endl;
    os << "  | History total allocated count:" << it.second.history_count
       << std::endl;
  }
  os << "  >=========================Current Heap Object========================<"
     << std::endl;
  for (const auto &obj : snapshot.mem_info_map_) {
    auto it = snapshot.hash_name_map_.find(obj.second.count_info_key);
    if (it == snapshot.hash_name_map_.end()) {
      continue;
    }
    os << "  +------------------------------+----------------------------"
       << std::endl;
    os << "  | Memory address:" << obj.first << std::endl;
    os << "  | Type:" << it->second.name << std::endl;
    os << "  | Source:" << obj.second.file_name << ":" << obj.second.line
       << std::endl;
  }
  os << "  >============================" + name +
            " End ============================<"
     << std::endl;
#endif // defined(DEBUG) || defined(_DEBUG)
}

auto kratos::service::MemoryAllocatorImpl::dump(std::ostream &os) -> void {
#if defined(DEBUG) || defined(_DEBUG)
  int num = 1;
  for (const auto &snapshot_info : snapshot_list_) {
    print_snapshot("Snapshot #" + std::to_string(num++), os,
                   snapshot_info.snapshot, snapshot_info.timestamp);
  }
  print_snapshot("Current", os, current_snapshot_,
                 util::get_os_time_millionsecond());
  return;
#endif // defined(DEBUG) || defined(_DEBUG)
}

auto kratos::service::MemoryAllocatorImpl::snapshot() -> void {
#if defined(DEBUG) || defined(_DEBUG)
  SnapshotInfo snapshot = {util::get_os_time_millionsecond(),
                           current_snapshot_};
  snapshot_list_.emplace_back(std::move(snapshot));
  return;
#endif // defined(DEBUG) || defined(_DEBUG)
}

auto kratos::service::MemoryAllocatorImpl::dump(int level,
                                                klogger::Appender *appender)
    -> void {
#if defined(DEBUG) || defined(_DEBUG)
  std::stringstream ss;
  int num = 1;
  for (const auto &snapshot_info : snapshot_list_) {
    print_snapshot("Snapshot #" + std::to_string(num++), ss,
                   snapshot_info.snapshot, snapshot_info.timestamp);
  }
  print_snapshot("Current", ss, current_snapshot_,
                 util::get_os_time_millionsecond());
  appender->write(level, "%s", ss.str().c_str());
#endif // defined(DEBUG) || defined(_DEBUG)
}
